package com.example.exercise.common.mvc;

import com.alibaba.fastjson.JSONObject;
import com.example.exercise.common.annotation.AccessLimit;
import com.example.exercise.common.base.BaseResponse;
import com.example.exercise.common.net.WebUtils;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;
import org.springframework.web.servlet.resource.ResourceHttpRequestHandler;

import java.io.OutputStreamWriter;
import java.util.concurrent.TimeUnit;

/**
 * 防止恶意刷单个接口的拦截器
 */
@Component
public class RateLimitApiInterceptor implements HandlerInterceptor {

    private static final int defaultPermitCount = 2;
    private static final int defaultDurationSecond = 1;

    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    /**
     * seconds秒内调接口超过count次就限流
     */
    @Override
    public boolean preHandle(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o) throws Exception {
        if(o instanceof ResourceHttpRequestHandler) {
            //资源拦截跳过
            return true;
        }
        HandlerMethod method = (HandlerMethod)o;
        int permitCount = defaultPermitCount;
        if(method.hasMethodAnnotation(AccessLimit.class)) {
            AccessLimit accessLimit = method.getMethodAnnotation(AccessLimit.class);
            permitCount = accessLimit.count();
        }
        String ip = WebUtils.getRemoteAddr(httpServletRequest);
        String key = "U_P_L:" + httpServletRequest.getServletPath() + ":" + ip;
        String string = redisTemplate.opsForValue().get(key);
        Integer count = Integer.valueOf(string == null? "0":string);
        if(count > permitCount){
            BaseResponse<String> baseResponse = new BaseResponse<>();
            baseResponse.errMsg("请求过于频繁，请" + permitCount + "秒后再试");
            OutputStreamWriter outputStreamWriter = new OutputStreamWriter(httpServletResponse.getOutputStream());
            outputStreamWriter.write(JSONObject.toJSONString(baseResponse));
            outputStreamWriter.flush();
            outputStreamWriter.close();
            return false;
        }
        return true;
    }

    @Override
    public void postHandle(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, ModelAndView modelAndView) {

    }

    @Override
    public void afterCompletion(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, Exception e) {
        if(o instanceof ResourceHttpRequestHandler) {
            //资源拦截跳过
            return;
        }
        HandlerMethod method = (HandlerMethod)o;
        int permitSeconds = defaultDurationSecond;
        if(method.hasMethodAnnotation(AccessLimit.class)) {
            AccessLimit accessLimit = method.getMethodAnnotation(AccessLimit.class);
            permitSeconds = accessLimit.seconds();
        }
        String ip = WebUtils.getRemoteAddr(httpServletRequest);
        String key = "U_P_L:" + httpServletRequest.getServletPath() + ":" + ip;
        redisTemplate.opsForValue().increment(key,1);
        redisTemplate.expire(key, permitSeconds, TimeUnit.SECONDS);
    }
}
